from time import time
from brian import *
from sys import argv
from pickle import load,dump

filename=argv[1]
myfile=open(filename,'rb')
x=load(myfile)
myfile.close()

dt = defaultclock.dt = 1*ms

#### equations
izh_eqs='''
dv/dt = (0.04/ms/mV)*v**2+(5/ms)*v+140*mV/ms-u   : volt
du/dt = a*(b*v-u)                                : volt/second
a                                                : 1/second
b                                                : 1/second
c                                                : volt
d                                               : volt/second
I                                               : volt/second
Fake                                            : volt/second
'''
reset_eqs='''
v=c;u+=d
'''

#threshold = 30*mV
threshold = x[0]*mV

def Izhikevich_scheme(G):
  for i in range(1):
    for i in range(2):
      G.v += 0.5*dt*((0.04/ms/mV)*G.v**2+(5.0/ms)*G.v+140*mV/ms-G.u+G.I)
    G.u += dt*G.a*(G.b*G.v-G.u)

def Izhikevich_scheme_rk4(G):
  for i in range(10):
    v1=(0.04/ms/mV)*G.v**2+(5.0/ms)*G.v+140*mV/ms-G.u+G.I
    u1=G.a*(G.b*G.v-G.u)
  
    v2=(0.04/ms/mV)*(G.v+0.5*0.1*dt*v1)**2+(5.0/ms)*(G.v+0.5*0.1*dt*v1)+140*mV/ms-(G.u+0.5*0.1*dt*u1)+G.I
    u2=G.a*(G.b*(G.v+0.5*0.1*dt*v1)-(G.u+0.5*0.1*dt*u1))
    
    v3=(0.04/ms/mV)*(G.v+0.5*0.1*dt*v2)**2+(5.0/ms)*(G.v+0.5*0.1*dt*v2)+140*mV/ms-(G.u+0.5*0.1*dt*u2)+G.I
    u3=G.a*(G.b*(G.v+0.5*0.1*dt*v2)-(G.u+0.5*0.1*dt*u2))
  
    v4=(0.04/ms/mV)*(G.v+0.1*dt*v3)**2+(5.0/ms)*(G.v+0.1*dt*v3)+140*mV/ms-(G.u+0.1*dt*u3)+G.I
    u4=G.a*(G.b*(G.v+0.1*dt*v3)-(G.u+0.1*dt*u3))
  
    G.v+=0.1666666667*0.1*dt*(v1+0.5*v2+0.5*v3+v4)
    G.u+=0.1666666667*0.1*dt*(u1+0.5*u2+0.5*u3+u4)
  pass



# G the neuron group
# N the number of neurons
# b the b constant
def make_Izhikevich_scheme(G,N,b):
    c1 = 0.04/ms/mV
    c2 = 5/ms
    c3 = 140*mV/ms
    # Izhikevich's numerical integration scheme
    def Izhikevich_scheme(G):
        G.v += (0.5*dt)*(c1*G.v**2+c2*G.v+c3-G.u+G.I)
        G.v += (0.5*dt)*(c1*G.v**2+c2*G.v+c3-G.u+G.I)
        G.u += dt*G.a*(b*G.v-G.u)
    weave_code = '''
    for(int i=0; i<N; i++){
        v[i] += (0.5*dt)*(c1*v[i]*v[i]+c2*v[i]+c3-u[i]+I[i]);
        v[i] += (0.5*dt)*(c1*v[i]*v[i]+c2*v[i]+c3-u[i]+I[i]);
        u[i] += dt*a[i]*(b*v[i]-u[i]);
    }
    '''
    weave_vars = (G.a, G.v, G.u, G.I, c1, c2, c3)
    weave_compiler = get_global_preference('weavecompiler')
    def weave_Izhikevich_scheme(G):
        a, v, u, I, c1, c2, c3 = weave_vars
        weave.inline(weave_code, ['a', 'v', 'u', 'I', 'c1', 'c2', 'c3', 'b', 'dt', 'N'],
                     compiler=weave_compiler, extra_compile_args=['-O3'])
    if get_global_preference('useweave'):
        return weave_Izhikevich_scheme
    return Izhikevich_scheme



#### General parameters
M=100.0 #number of synapse per neuron
D=1 #maximal conduction delay
Ne=800 #number of excitatory neurons
Ni=200 #number of inhibitory neurons
N=Ne+Ni
sm=26.0*mV/ms #maximum synaptic strength
#delay_pre_reinf=rand()*2000*ms
interval_pre_post=40*ms

n1=12
n2=14
n3=16
STDP_tau_pre= 20*ms #
STDP_coef_pre = 1.0
STDP_tau_post= 20*ms #
STDP_coef_post = -1.5

DA_tau=200*ms
display=True
DA_burst_val=0.1e-4
#### variables

re=rand(Ne)
ri=rand(Ni)

a_e=0.02/ms
b_e=0.2/ms
c_e=-65.0*mV
d_e=8.0*mV/ms

a_i=0.1/ms
b_i=0.2/ms
c_i=-65.0*mV
d_i=2.0*mV/ms


###### neuron group
neurons = NeuronGroup(1000, model=izh_eqs, threshold=threshold, reset=reset_eqs)

# Izhikevich's numerical integration scheme
#neurons._state_updater = Izhikevich_scheme
neurons._state_updater = make_Izhikevich_scheme( neurons , N, b_e )

sub_e=neurons.subgroup(Ne)
sub_i=neurons.subgroup(Ni)

sub_e.a=a_e
sub_e.b=b_e
sub_e.c=c_e
sub_e.d=d_e

sub_i.a=a_i
sub_i.b=b_i
sub_i.c=c_i
sub_i.d=d_i


linkse=DelayConnection(sub_e,neurons,'I',structure='sparse',sparseness=M/1000.0,weight=1.0*mV/ms,max_delay=100*ms,delay=(0*ms, 100*ms))
linkse.W[n1,n2]=2.0*mV
linkse.W[n1,n3]=2.0*mV
linkse_t=DelayConnection(sub_e,neurons,'Fake',structure='sparse',weight=linkse.W,max_delay=100*ms)
#linkse_t.set_delays(linkse.delay) #doesn't work for now, set delays elsewhere
forget(linkse_t)
stdpe = ExponentialSTDP( linkse_t, STDP_tau_pre, STDP_tau_post, STDP_coef_pre, STDP_coef_post, interactions='all',wmin=-1, wmax=1,update='additive')
linkse_t.compress()
linkse_t.delay.alldata=linkse_t.delay.alldata#initialize the trace weights correctly
linkse_t.W.alldata*=0.0001#initialize the trace weights correctly

linksi=DelayConnection(sub_i,sub_e,'I',structure='sparse',sparseness=M/800.0,weight=-1.0*mV/ms,max_delay=10*ms,delay=(0*ms, 100*ms))
#linksi_t=DelayConnection(sub_i,sub_e,'Fake',structure='sparse',weight=linksi.W,max_delay=100*ms)
#linksi_t.set_delays(linksi.delay)
#forget(linksi_t)
#stdpi = ExponentialSTDP( linksi_t, STDP_tau_pre, STDP_tau_post, STDP_coef_pre, STDP_coef_post, interactions='all',wmin=-10, wmax=10,update='additive')
#linksi_t.compress()
#linksi_t.delay.alldata=linksi_t.delay.alldata#initialize the trace weights correctly
#linksi_t.W.alldata*=0.0001#initialize the trace weights correctly


#### Initial state
neurons.v=-65*mV
neurons.u=neurons.v*neurons.b

spkm_n=SpikeMonitor(neurons)
spkc_n=SpikeCounter(neurons)

#### reinforcement signal use
# logging
weightn1n2 = []
magic_register( weightn1n2 )
weightn1n3 = []
magic_register( weightn1n3 )

# reinforcement info
reinf = [] #reinforcement signal
magic_register( reinf )
reinf.append(0)
reinf_buffer = [] #spike times buffer
magic_register( reinf_buffer )
reinf_spk = [] #spike times
magic_register( reinf_spk )
n1_hist = [] #spike times of n1
magic_register( n1_hist )


def f(spikes):
  if len(reinf)==0: #first step
    reinf.append(0.0)
  else:
    reinf.append(reinf[-1]*0.995) # 1-1/200
  if n2 in spikes:
    #print "n2 spktime " + str(defaultclock.t)
    if len(n1_hist)>0 and n1_hist[-1]+interval_pre_post>=defaultclock.t: #if the last spike of n1 is from less than 20*ms
        reinf_buffer.append([defaultclock.t+2000*ms*rand(),DA_burst_val])#add to the reinforcements to be included
        #print "reinf pos added to be done later "+ str(reinf_buffer[-1][0])
  if n3 in spikes:
    #print "n3 spktime " + str(defaultclock.t)
    if len(n1_hist)>0 and n1_hist[-1]+interval_pre_post>=defaultclock.t: #if the last spike of n1 is from less than 20*ms
        reinf_buffer.append([defaultclock.t+2000*ms*rand(),-DA_burst_val])#add to the reinforcements to be included
        #print "reinf neg added to be done later "+ str(reinf_buffer[-1][0])
  if n1 in spikes:
    #print "n1 spktime "+ str(defaultclock.t)
    n1_hist.append(defaultclock.t)

  for reinf_time in reinf_buffer:
    if abs(defaultclock.t-reinf_time[0])<1*ms and defaultclock.t-reinf_time[0]>=0:
      #print "reinf spktime "+ str(defaultclock.t)
      reinf[-1]+=reinf_time[1]
      reinf_spk.append(reinf_time)
      reinf_buffer.remove(reinf_time)
  weightn1n2.append(linkse.W[n1,n2])
  weightn1n3.append(linkse.W[n1,n3])

reinf_spk_mon=SpikeMonitor(neurons,function=f)
#### runtime operations


@network_operation(clock=EventClock(dt=10*ms))
def trace_reduction():
  if abs(reinf[-1])>1e-10:#weight update if necessary
    linkse.W.alldata+=(reinf[-1])*(linkse_t.W.alldata)*50.0
    linkse.W.alldata[:]=clip(linkse.W.alldata, 0,sm)
  linkse_t.W.alldata*=0.99 
#  linksi_t.W.alldata*=0.99 

@network_operation(when='before_connections')
def random_thalamic_input():
    neurons.I=(13*randn(N)-0.5)*mV/ms

net=Network(neurons,linkse,linkse_t,linksi,spkm_n,spkc_n,trace_reduction,random_thalamic_input,reinf_spk_mon)
##### run
seconds=5
net.run(seconds*1000*ms)
#print str(spkc_n.nspikes)
#raster_plot()
plot(range(len(weightn1n2)),weightn1n2,'b')
plot(range(len(weightn1n3)),weightn1n3,'g')
assert len( weightn1n2 ) > 0
assert len( weightn1n3 ) > 0 

myfile=open(filename,'wb')
dump([weightn1n2[-1]-weightn1n3[-1]], myfile)
myfile.close()

